# Copyright (c) OpenMMLab. All rights reserved.
from typing import List, Union

import torch.nn as nn
from mmdet.utils import ConfigType, OptMultiConfig

from mmyolo.registry import MODELS
from mmcv.cnn import ConvModule
from .. import CSPLayerWithTwoConv
from ..utils import make_divisible, make_round
from .yolov5_pafpn import YOLOv5PAFPN
from .base_yolo_neck import BaseYOLONeck
from ..layers import (RepNCSPELAN4, SPPELAN, ADown, 
                     CBLinear, CBFuse, Concat, Conv, DFL, RepConvN)


@MODELS.register_module()
class YOLOv8PAFPN(YOLOv5PAFPN):
    print(" THIS IS IN THE BASE YOLOv8PAFPN REPLACE SPP")
    # This is the original 
    def __init__(self,
                 in_channels: List[int],
                 out_channels: Union[List[int], int],
                 deepen_factor: float = 1.25,
                 widen_factor: float = 1.25,
                 num_csp_blocks: int = 3,
                 freeze_all: bool = False,
                 norm_cfg: ConfigType = dict(
                     type='BN', momentum=0.03, eps=0.001),
                 act_cfg: ConfigType = dict(type='SiLU', inplace=True),
                 init_cfg: OptMultiConfig = None):
        super().__init__(
            in_channels=in_channels,
            out_channels=out_channels,
            deepen_factor=deepen_factor,
            widen_factor=widen_factor,
            num_csp_blocks=num_csp_blocks,
            freeze_all=freeze_all,
            norm_cfg=norm_cfg,
            act_cfg=act_cfg,
            init_cfg=init_cfg)

    def build_reduce_layer(self, idx: int) -> nn.Module:
        """build reduce layer.
        Args:
            idx (int): layer idx.
        Returns:
            nn.Module: The reduce layer.
        """
        return nn.Identity()

    # def build_top_down_layer(self, idx: int) -> nn.Module:
    #     """build top down layer.
    #     Args:
    #         idx (int): layer idx.
    #     Returns:
    #         nn.Module: The top down layer.
    #     """
    #     return CSPLayerWithTwoConv(
    #         make_divisible((self.in_channels[idx - 1] + self.in_channels[idx]),
    #                        self.widen_factor),
    #         make_divisible(self.out_channels[idx - 1], self.widen_factor),
    #         num_blocks=make_round(self.num_csp_blocks, self.deepen_factor),
    #         add_identity=False,
    #         norm_cfg=self.norm_cfg,
    #         act_cfg=self.act_cfg)

    def build_top_down_layer(self, idx: int) -> nn.Module:
        """Build ELAN-based top-down layer with SPP integration."""
        in_channels_combined = make_divisible((self.in_channels[idx - 1] + self.in_channels[idx]), 
        self.widen_factor)
        out_channels_prev =  make_divisible(self.out_channels[idx - 1], self.widen_factor)
        hidden =  make_divisible(out_channels_prev * 0.5 , self.widen_factor)
        
        return RepNCSPELAN4(
            in_channels_combined,
            out_channels_prev,
            hidden * 2,
            hidden,
            make_round(self.num_csp_blocks, self.widen_factor)
        )

    def build_bottom_up_layer(self, idx: int) -> nn.Module:
        """build bottom up layer.
        Args:
            idx (int): layer idx.
        Returns:
            nn.Module: The bottom up layer.
        """
    #     return CSPLayerWithTwoConv(
    #         make_divisible((self.out_channels[idx] + self.out_channels[idx + 1]),self.widen_factor),
    #         make_divisible(self.out_channels[idx + 1], self.widen_factor),
    #         num_blocks=make_round(self.num_csp_blocks, self.deepen_factor),
    #         add_identity=False,
    #         norm_cfg=self.norm_cfg,
    #         act_cfg=self.act_cfg)
        
        in_channels_combined = make_divisible((self.out_channels[idx] + self.out_channels[idx + 1]),self.widen_factor)
        out_channels_prev =  make_divisible(self.out_channels[idx + 1], self.widen_factor)
        hidden =  make_divisible(out_channels_prev * 0.5 , self.widen_factor)

        return RepNCSPELAN4(
            in_channels_combined,
            out_channels_prev,
            hidden * 2,
            hidden,
            make_round(self.num_csp_blocks, self.widen_factor)
        )
